# -*- coding: utf-8 -*-
"""
PostGIS数据导入器
根据GeoJSON和GeoPackage处理器返回的几何类型，将数据导入到对应的PostGIS表中
"""

import os
import sys
from typing import Dict, Any, Optional
from sqlalchemy import create_engine, inspect, text
from sqlalchemy.exc import SQLAlchemyError
from loguru import logger

# 添加处理器模块路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from hydropulse.utils.geojson_processor import GeoJSONProcessor
from hydropulse.utils.geopackage_processor import GeoPackageProcessor
from hydropulse.utils.shp_processor import SHPProcessor


class PostGISImporter:
    """PostGIS数据导入器"""
    
    def __init__(self, db_url: str):
        """
        初始化PostGIS导入器
        
        Args:
            db_url: 数据库连接字符串，格式：postgresql://user:password@host:port/database
        """
        self.db_url = db_url
        self.engine = create_engine(db_url)
        
    def _get_target_table(self, geometry_type: str) -> str:
        """
        根据几何类型获取目标表名
        
        Args:
            geometry_type: 几何类型字符串
            
        Returns:
            目标表名
        """
        geometry_type = geometry_type.upper()
        
        # 映射几何类型到对应表
        if geometry_type in ['POINT', 'MULTIPOINT']:
            return 'geometry_points'
        elif geometry_type in ['POLYGON', 'MULTIPOLYGON']:
            return 'geometry_multipolygons'
        else:
            logger.warning(f"未知几何类型 '{geometry_type}'，默认使用 geometry_multipolygons")
            return 'geometry_multipolygons'
    
    def _check_table_exists(self, table_name: str) -> bool:
        """
        检查表是否存在
        
        Args:
            table_name: 表名
            
        Returns:
            表是否存在
        """
        inspector = inspect(self.engine)
        return inspector.has_table(table_name)
    
    def import_geojson(self, file_path: str, table_name: Optional[str] = None) -> Dict[str, Any]:
        """
        导入GeoJSON文件到PostGIS
        
        Args:
            file_path: GeoJSON文件路径(.json/.geojson)
            table_name: 可选，指定表名，如果不提供则根据几何类型自动选择
            
        Returns:
            导入结果字典
        """
        try:
            logger.info(f"开始导入GeoJSON文件: {file_path}")
            
            # 使用GeoJSON处理器处理数据
            from .geojson_processor import GeoJSONProcessor
            processor = GeoJSONProcessor(file_path)
            result = processor.load_data()
            
            if not result or 'data' not in result:
                raise ValueError("无法加载GeoJSON数据")
            
            gdf = result['data']
            
            # 验证数据
            if gdf.empty:
                raise ValueError("GeoJSON不包含任何数据")
            
            # 获取几何类型
            geometry_type = result['geometry_type']
            
            # 确定目标表
            if table_name is None:
                table_name = self._get_target_table(geometry_type)
            
            # 检查表是否存在
            if not self._check_table_exists(table_name):
                logger.error(f"表 '{table_name}' 不存在，请先创建表")
                return {
                    'success': False,
                    'error': f"表 '{table_name}' 不存在，请先创建表",
                    'file_path': file_path
                }
            
            # 处理重复键问题：跳过已存在的记录
            existing_codes = set()
            try:
                with self.engine.connect() as conn:
                    result = conn.execute(text(f"SELECT code FROM {table_name}"))
                    existing_codes = {row[0] for row in result}
            except Exception as e:
                logger.warning(f"无法获取现有code列表: {e}")
            
            # 过滤掉已存在的记录
            if 'code' in gdf.columns and existing_codes:
                original_count = len(gdf)
                gdf = gdf[~gdf['code'].isin(existing_codes)]
                filtered_count = original_count - len(gdf)
                if filtered_count > 0:
                    logger.info(f"跳过 {filtered_count} 条已存在的记录")
            
            # 如果没有新数据需要导入
            if gdf.empty:
                logger.info("所有记录都已存在，无需导入")
                return {
                    'success': True,
                    'table_name': table_name,
                    'record_count': 0,
                    'skipped_count': original_count if 'original_count' in locals() else 0,
                    'geometry_type': geometry_type,
                    'file_path': file_path,
                    'message': '所有记录都已存在，无需导入'
                }
            
            # 导入新数据
            gdf.to_postgis(table_name, self.engine, if_exists='append', index=False)
            
            logger.info(f"成功导入 {len(gdf)} 条新记录到表 '{table_name}'")
            
            return {
                'success': True,
                'table_name': table_name,
                'record_count': len(gdf),
                'skipped_count': original_count - len(gdf) if 'original_count' in locals() else 0,
                'geometry_type': geometry_type,
                'file_path': file_path
            }
            
        except Exception as e:
            logger.error(f"导入GeoJSON失败: {e}")
            return {
                'success': False,
                'error': str(e),
                'file_path': file_path
            }
    
    def import_geopackage(self, file_path: str, layer_name: Optional[str] = None, table_name: Optional[str] = None) -> Dict[str, Any]:
        """
        导入GeoPackage文件到PostGIS
        
        Args:
            file_path: GeoPackage文件路径(.gpkg)
            layer_name: 可选，指定图层名，如果不提供则导入第一个图层
            table_name: 可选，指定表名，如果不提供则根据几何类型自动选择
            
        Returns:
            导入结果字典
        """
        try:
            logger.info(f"开始导入GeoPackage文件: {file_path}")
            
            # 使用GeoPackage处理器处理数据
            from .geopackage_processor import GeoPackageProcessor
            processor = GeoPackageProcessor(file_path)
            result = processor.load_data(layer_name=layer_name)
            
            if not result or 'data' not in result:
                raise ValueError("无法加载GeoPackage数据")
            
            gdf = result['data']
            
            # 验证数据
            if gdf.empty:
                raise ValueError("GeoPackage不包含任何数据")
            
            # 获取几何类型
            geometry_type = result['geometry_type']
            
            # 确定目标表
            if table_name is None:
                table_name = self._get_target_table(geometry_type)
            
            # 检查表是否存在
            if not self._check_table_exists(table_name):
                logger.error(f"表 '{table_name}' 不存在，请先创建表")
                return {
                    'success': False,
                    'error': f"表 '{table_name}' 不存在，请先创建表",
                    'file_path': file_path
                }
            
            # 处理重复键问题：跳过已存在的记录
            existing_codes = set()
            try:
                with self.engine.connect() as conn:
                    result = conn.execute(text(f"SELECT code FROM {table_name}"))
                    existing_codes = {row[0] for row in result}
            except Exception as e:
                logger.warning(f"无法获取现有code列表: {e}")
            
            # 过滤掉已存在的记录
            if 'code' in gdf.columns and existing_codes:
                original_count = len(gdf)
                gdf = gdf[~gdf['code'].isin(existing_codes)]
                filtered_count = original_count - len(gdf)
                if filtered_count > 0:
                    logger.info(f"跳过 {filtered_count} 条已存在的记录")
            
            # 如果没有新数据需要导入
            if gdf.empty:
                logger.info("所有记录都已存在，无需导入")
                return {
                    'success': True,
                    'table_name': table_name,
                    'record_count': 0,
                    'skipped_count': original_count if 'original_count' in locals() else 0,
                    'geometry_type': geometry_type,
                    'layer_name': layer_name or 'default',
                    'file_path': file_path,
                    'message': '所有记录都已存在，无需导入'
                }
            
            # 导入新数据
            gdf.to_postgis(table_name, self.engine, if_exists='append', index=False)
            
            logger.info(f"成功导入 {len(gdf)} 条新记录到表 '{table_name}'")
            
            return {
                'success': True,
                'table_name': table_name,
                'record_count': len(gdf),
                'skipped_count': original_count - len(gdf) if 'original_count' in locals() else 0,
                'geometry_type': geometry_type,
                'layer_name': layer_name or 'default',
                'file_path': file_path
            }
            
        except Exception as e:
            logger.error(f"导入GeoPackage失败: {e}")
            return {
                'success': False,
                'error': str(e),
                'file_path': file_path
            }
    
    def import_file(self, file_path: str, **kwargs) -> Dict[str, Any]:
        """
        根据文件类型自动选择导入方法
        
        Args:
            file_path: 文件路径
            **kwargs: 其他参数传递给具体的导入方法
            
        Returns:
            导入结果字典
        """
        file_path = file_path.lower()
        
        if file_path.endswith('.geojson') or file_path.endswith('.json'):
            return self.import_geojson(file_path, **kwargs)
        elif file_path.endswith('.gpkg'):
            return self.import_geopackage(file_path, **kwargs)
        elif file_path.endswith('.shp'):
            return self.import_shapefile(file_path, **kwargs)
        else:
            return {
                'success': False,
                'error': f"不支持的文件格式: {file_path}",
                'file_path': file_path
            }

    def import_shapefile(self, file_path: str, table_name: Optional[str] = None) -> Dict[str, Any]:
        """
        导入Shapefile文件到PostGIS
        
        Args:
            file_path: Shapefile文件路径(.shp)
            table_name: 可选，指定表名，如果不提供则根据几何类型自动选择
            
        Returns:
            导入结果字典
        """
        try:
            logger.info(f"开始导入Shapefile文件: {file_path}")
            
            # 使用SHP处理器处理数据
            from .shp_processor import SHPProcessor
            processor = SHPProcessor(file_path)
            result = processor.load_data()
            
            if not result or 'data' not in result:
                raise ValueError("无法加载Shapefile数据")
            
            gdf = result['data']
            
            # 验证数据
            if gdf.empty:
                raise ValueError("Shapefile不包含任何数据")
            
            # 获取几何类型
            geometry_type = result['geometry_type']
            
            # 确定目标表
            if table_name is None:
                table_name = self._get_target_table(geometry_type)
            
            # 检查表是否存在
            if not self._check_table_exists(table_name):
                logger.error(f"表 '{table_name}' 不存在，请先创建表")
                return {
                    'success': False,
                    'error': f"表 '{table_name}' 不存在，请先创建表",
                    'file_path': file_path
                }
            
            # 处理重复键问题：跳过已存在的记录
            existing_codes = set()
            try:
                with self.engine.connect() as conn:
                    result = conn.execute(text(f"SELECT code FROM {table_name}"))
                    existing_codes = {row[0] for row in result}
            except Exception as e:
                logger.warning(f"无法获取现有code列表: {e}")
            
            # 过滤掉已存在的记录
            if 'code' in gdf.columns and existing_codes:
                original_count = len(gdf)
                gdf = gdf[~gdf['code'].isin(existing_codes)]
                filtered_count = original_count - len(gdf)
                if filtered_count > 0:
                    logger.info(f"跳过 {filtered_count} 条已存在的记录")
            
            # 如果没有新数据需要导入
            if gdf.empty:
                logger.info("所有记录都已存在，无需导入")
                return {
                    'success': True,
                    'table_name': table_name,
                    'record_count': 0,
                    'skipped_count': original_count if 'original_count' in locals() else 0,
                    'geometry_type': geometry_type,
                    'file_path': file_path,
                    'message': '所有记录都已存在，无需导入'
                }
            
            # 导入新数据
            gdf.to_postgis(table_name, self.engine, if_exists='append', index=False)
            
            logger.info(f"成功导入 {len(gdf)} 条新记录到表 '{table_name}'")
            
            return {
                'success': True,
                'table_name': table_name,
                'record_count': len(gdf),
                'skipped_count': original_count - len(gdf) if 'original_count' in locals() else 0,
                'geometry_type': geometry_type,
                'file_path': file_path
            }
            
        except Exception as e:
            logger.error(f"导入Shapefile失败: {e}")
            return {
                'success': False,
                'error': str(e),
                'file_path': file_path
            }
    
    def get_table_info(self, table_name: str) -> Dict[str, Any]:
        """
        获取表信息
        
        Args:
            table_name: 表名
            
        Returns:
            表信息字典
        """
        try:
            with self.engine.connect() as conn:
                # 获取记录数
                result = conn.execute(text(f"SELECT COUNT(*) FROM {table_name}"))
                count = result.scalar()
                
                # 获取几何类型
                result = conn.execute(text(f"SELECT ST_GeometryType(geometry) FROM {table_name} LIMIT 1"))
                geom_type = result.scalar()
                
                return {
                    'table_name': table_name,
                    'record_count': count,
                    'geometry_type': geom_type
                }
                
        except Exception as e:
            logger.error(f"获取表信息失败: {e}")
            return {
                'error': str(e),
                'table_name': table_name
            }
    
    def clear_table(self, table_name: str) -> Dict[str, Any]:
        """
        清空指定表中的数据
        
        Args:
            table_name: 要清空的表名
            
        Returns:
            操作结果字典
        """
        try:
            if not self._check_table_exists(table_name):
                logger.error(f"表 '{table_name}' 不存在")
                return {
                    'success': False,
                    'error': f"表 '{table_name}' 不存在",
                    'table_name': table_name
                }
            
            with self.engine.connect() as conn:
                # 获取清空前的记录数
                count_before = conn.execute(text(f"SELECT COUNT(*) FROM {table_name}")).scalar()
                
                # 清空表数据
                conn.execute(text(f"DELETE FROM {table_name}"))
                conn.commit()
                
                # 获取清空后的记录数
                count_after = conn.execute(text(f"SELECT COUNT(*) FROM {table_name}")).scalar()
                
                logger.info(f"成功清空表 '{table_name}'，删除 {count_before} 条记录")
                
                return {
                    'success': True,
                    'table_name': table_name,
                    'deleted_count': count_before,
                    'remaining_count': count_after
                }
                
        except Exception as e:
            logger.error(f"清空表失败: {e}")
            return {
                'success': False,
                'error': str(e),
                'table_name': table_name
            }
    
    def clear_tables(self, table_names: list) -> Dict[str, Any]:
        """
        批量清空多个表中的数据
        
        Args:
            table_names: 要清空的表名列表
            
        Returns:
            操作结果字典
        """
        results = []
        total_deleted = 0
        
        for table_name in table_names:
            result = self.clear_table(table_name)
            results.append(result)
            if result['success']:
                total_deleted += result['deleted_count']
        
        return {
            'success': True,
            'results': results,
            'total_deleted': total_deleted,
            'tables_processed': len(table_names)
        }


if __name__ == "__main__":
    # 示例用法
    import argparse
    
    parser = argparse.ArgumentParser(description='PostGIS数据导入器')
    parser.add_argument('--db-url', required=True, help='数据库连接字符串')
    parser.add_argument('--file', required=True, help='要导入的文件路径')
    parser.add_argument('--layer', help='GeoPackage图层名称（可选）')
    parser.add_argument('--table', help='目标表名（可选）')
    
    args = parser.parse_args()
    
    # 创建导入器
    importer = PostGISImporter(args.db_url)
    
    # 导入文件
    result = importer.import_file(args.file, layer_name=args.layer, table_name=args.table)
    
    if result['success']:
        print(f"✅ 导入成功: {result['record_count']} 条记录导入到表 '{result['table_name']}'")
    else:
        print(f"❌ 导入失败: {result['error']}")
    
    # 测试代码
    """
    # 使用示例
    db_url = "postgresql://username:password@localhost:5432/database"
    importer = PostGISImporter(db_url)
    
    # 导入GeoJSON
    result1 = importer.import_geojson("碧流河.json")
    
    # 导入GeoPackage
    result2 = importer.import_geopackage("gansu.gpkg")
    
    # 导入Shapefile
    result3 = importer.import_shapefile("水库合并shp/merged_output.shp")
    
    # 自动识别文件类型
    result4 = importer.import_file("test_data.geojson")
    result5 = importer.import_file("test_data.shp")
    """